package com.liuhengjia.langchain4japply.api;

import com.liuhengjia.langchain4japply.config.assistant.Assistant;
import com.liuhengjia.langchain4japply.config.assistant.MemoryIsolationAssistant;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import lombok.RequiredArgsConstructor;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 会话隔离 API
 */
@RestController
@RequestMapping("/api/memory-isolation")
@RequiredArgsConstructor
public class MemoryIsolationAPI {
    private final ChatLanguageModel chatLanguageModel;
    private final Assistant memoryIsolationAssistant;
    private static final Map<String, ChatMemory> STRING_CHAT_MEMORY_MAP = new ConcurrentHashMap<>();

    /**
     * 低级聊天
     *
     * @param memoryId 会话 ID
     * @param message  用户信息
     * @return 模型响应
     */
    @GetMapping("/low/chat")
    public String selfChat(@RequestParam(value = "memoryId") String memoryId, @RequestParam(value = "message") String message) {
        if (!STRING_CHAT_MEMORY_MAP.containsKey(memoryId)) {
            STRING_CHAT_MEMORY_MAP.put(memoryId, MessageWindowChatMemory.builder().maxMessages(10).build());
        }
        ChatMemory currentChatMemory = STRING_CHAT_MEMORY_MAP.get(memoryId);
        currentChatMemory.add(UserMessage.from(message));
        AiMessage aiMessage = chatLanguageModel.chat(currentChatMemory.messages()).aiMessage();
        currentChatMemory.add(aiMessage);
        return aiMessage.text();
    }

    /**
     * 高级聊天
     *
     * @param memoryId 会话 ID
     * @param message  用户信息
     * @return 模型响应
     */
    @GetMapping("/high/chat")
    public String highChat(@RequestParam(value = "memoryId") String memoryId, @RequestParam(value = "message") String message) {
        return ((MemoryIsolationAssistant) memoryIsolationAssistant).memoryIsolationChat(memoryId, message);
    }
}
